Update app.py

app.py

@@ -38,25 +38,33 @@ if api_key:
     CLIENT = setup_roboflow_client(api_key)
 
     def segment_image(image_path):
-    results = CLIENT.infer(image_path, model_id="closet/1")
-    # json.loads() 제거, results를 직접 사용
-    
-    image = cv2.imread(image_path)
-    image = cv2.resize(image, (800, 600))
-    mask = np.zeros(image.shape, dtype=np.uint8)
-    
-    for prediction in results['predictions']:
-        points = prediction['points']
-        pts = np.array([[p['x'], p['y']] for p in points], np.int32)
-        scale_x = image.shape[1] / results['image']['width']
-        scale_y = image.shape[0] / results['image']['height']
-        pts = pts * [scale_x, scale_y]
-        pts = pts.astype(np.int32)
-        pts = pts.reshape((-1, 1, 2))
-        cv2.fillPoly(mask, [pts], color=(255, 255, 255))  # White mask
-    
-    segmented_image = cv2.bitwise_and(image, mask)
-    return Image.fromarray(cv2.cvtColor(segmented_image, cv2.COLOR_BGR2RGB))
+        try:
+            results = CLIENT.infer(image_path, model_id="closet/1")
+            
+            image = cv2.imread(image_path)
+            image = cv2.resize(image, (800, 600))
+            mask = np.zeros(image.shape, dtype=np.uint8)
+            
+            if 'predictions' in results:
+                for prediction in results['predictions']:
+                    points = prediction['points']
+                    pts = np.array([[p['x'], p['y']] for p in points], np.int32)
+                    scale_x = image.shape[1] / results['image']['width']
+                    scale_y = image.shape[0] / results['image']['height']
+                    pts = pts * [scale_x, scale_y]
+                    pts = pts.astype(np.int32)
+                    pts = pts.reshape((-1, 1, 2))
+                    cv2.fillPoly(mask, [pts], color=(255, 255, 255))  # White mask
+            
+                segmented_image = cv2.bitwise_and(image, mask)
+            else:
+                st.warning("No predictions found in the image. Returning original image.")
+                segmented_image = image
+            
+            return Image.fromarray(cv2.cvtColor(segmented_image, cv2.COLOR_BGR2RGB))
+        except Exception as e:
+            st.error(f"Error in segmentation: {str(e)}")
+            return Image.open(image_path)  # Return original image if segmentation fails
 
     def get_image_embedding(image):
         image_tensor = preprocess_val(image).unsqueeze(0).to(device)