Update app.py

app.py

@@ -148,23 +148,18 @@ class YOLOWorldDetector:
         
         print(f"Loading {self.model_name} on {self.device}...")
         try:
-            # Use the correct repository ID
-            self.model = AutoModel.from_pretrained(
-                f"IDEA-Research/{self.model_name}", 
-                trust_remote_code=True
-            )
-            self.model.to(self.device)
-            self.processor = AutoProcessor.from_pretrained(
-                f"IDEA-Research/{self.model_name}"
-            )
-            print("Model loaded successfully!")
+            # Try to load using Ultralytics YOLOWorld
+            from ultralytics import YOLOWorld
+            self.model = YOLOWorld(self.model_name)
+            self.model_type = "yoloworld"
+            print("YOLOWorld model loaded successfully!")
         except Exception as e:
-            print(f"Error loading model: {e}")
-            print("Falling back to YOLOv8 for detection...")
-            # Fallback to YOLOv8 if YOLOWorld fails to load
-            self.model = None
-            self.processor = None
-            self.fallback_model = YOLO("yolov8n.pt")
+            print(f"Error loading YOLOWorld model: {e}")
+            print("Falling back to standard YOLOv8 for detection...")
+            # Fallback to YOLOv8
+            self.model = YOLO("yolov8n.pt")
+            self.model_type = "yolov8"
+            print("YOLOv8 fallback model loaded successfully!")
         
         # Segmentation models
         self.seg_models = {}
@@ -176,15 +171,18 @@ class YOLOWorldDetector:
             
             print(f"Loading {self.model_name} on {self.device}...")
             try:
-                # Use Ultralytics YOLOWorld model
+                # Try to load using Ultralytics YOLOWorld
                 from ultralytics import YOLOWorld
                 self.model = YOLOWorld(self.model_name)
-                print("Model loaded successfully!")
+                self.model_type = "yoloworld"
+                print("YOLOWorld model loaded successfully!")
             except Exception as e:
                 print(f"Error loading YOLOWorld model: {e}")
                 print("Falling back to standard YOLOv8 for detection...")
-                # Fallback to YOLOv8 if YOLOWorld fails to load
+                # Fallback to YOLOv8
                 self.model = YOLO("yolov8n.pt")
+                self.model_type = "yolov8"
+                print("YOLOv8 fallback model loaded successfully!")
         return f"Using {self.model_name} model"
 
     def load_seg_model(self, model_name):
@@ -198,28 +196,33 @@ class YOLOWorldDetector:
         if image is None:
             return None, "No image provided"
         
-        try:
-            # Check if we're using YOLOWorld or standard YOLO
-            from ultralytics import YOLOWorld
-            is_yoloworld = isinstance(self.model, YOLOWorld)
-        except:
-            is_yoloworld = False
-            
         # Process the image
         if isinstance(image, str):
             img_for_json = cv2.imread(image)
         elif isinstance(image, np.ndarray):
             img_for_json = image.copy()
+        else:
+            # Convert PIL Image to numpy array if needed
+            img_for_json = np.array(image)
         
-        # Run inference
-        if is_yoloworld:
-            # YOLOWorld supports text prompts
-            results = self.model.predict(
-                source=image,
-                classes=text_prompt.split(','),
-                conf=confidence_threshold,
-                verbose=False
-            )
+        # Run inference based on model type
+        if self.model_type == "yoloworld":
+            try:
+                # YOLOWorld supports text prompts
+                results = self.model.predict(
+                    source=image,
+                    classes=text_prompt.split(','),
+                    conf=confidence_threshold,
+                    verbose=False
+                )
+            except Exception as e:
+                print(f"Error during YOLOWorld inference: {e}")
+                # If YOLOWorld inference fails, try to use it as standard YOLO
+                results = self.model.predict(
+                    source=image,
+                    conf=confidence_threshold,
+                    verbose=False
+                )
         else:
             # Standard YOLO doesn't use text prompts
             results = self.model.predict(