Update app.py

app.py

@@ -167,23 +167,23 @@ def process_image_detection(image, target_label, surprise_rating):
             image = image.convert('RGB')
 
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        print(f"Using device: {device}")  # Debug print
+        print(f"Using device: {device}")
 
         # Get original image DPI and size
         original_dpi = image.info.get('dpi', (72, 72))
         original_size = image.size
-        print(f"Image size: {original_size}")  # Debug print
+        print(f"Image size: {original_size}")
 
         # Calculate relative font size based on image dimensions
         base_fontsize = min(original_size) / 40
 
-        print("Loading models...")  # Debug print
+        print("Loading models...")
         owlv2_processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16")
         owlv2_model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16").to(device)
         sam_processor = AutoProcessor.from_pretrained("facebook/sam-vit-base")
         sam_model = AutoModelForMaskGeneration.from_pretrained("facebook/sam-vit-base").to(device)
 
-        print("Running object detection...")  # Debug print
+        print("Running object detection...")
         inputs = owlv2_processor(text=[target_label], images=image, return_tensors="pt").to(device)
         with torch.no_grad():
             outputs = owlv2_model(**inputs)
@@ -204,10 +204,10 @@ def process_image_detection(image, target_label, surprise_rating):
             max_score = scores[max_score_idx].item()
 
             if max_score > 0.2:
-                print("Processing detection results...")  # Debug print
+                print("Processing detection results...")
                 box = results["boxes"][max_score_idx].cpu().numpy()
 
-                print("Running SAM model...")  # Debug print
+                print("Running SAM model...")
                 # Convert image to numpy array if needed for SAM
                 if isinstance(image, Image.Image):
                     image_np = np.array(image)
@@ -215,7 +215,7 @@ def process_image_detection(image, target_label, surprise_rating):
                     image_np = image
 
                 sam_inputs = sam_processor(
-                    image_np,  # Use numpy array here
+                    image_np,
                     input_boxes=[[[box[0], box[1], box[2], box[3]]]],
                     return_tensors="pt"
                 ).to(device)
@@ -229,7 +229,7 @@ def process_image_detection(image, target_label, surprise_rating):
                     sam_inputs["reshaped_input_sizes"].cpu()
                 )
 
-                print(f"Mask type: {type(masks)}, Mask shape: {len(masks)}")  # Debug print
+                print(f"Mask type: {type(masks)}, Mask shape: {len(masks)}")
                 mask = masks[0]
                 if isinstance(mask, torch.Tensor):
                     mask = mask.numpy()
@@ -266,24 +266,23 @@ def process_image_detection(image, target_label, surprise_rating):
                 )
 
         plt.axis('off')
-    
-        print("Saving final image...")  # Debug print
+
+        print("Saving final image...")
         try:
-            # Save to buffer
-            buf = io.BytesIO()
-            
-            # Force figure to be in a format we can save
+            # Force figure to be rendered
             fig.canvas.draw()
             
-            # Get the image data from the figure
-            plot_data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
-            plot_data = plot_data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+            # Get the RGBA buffer from the figure
+            w, h = fig.canvas.get_width_height()
+            buf = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+            buf.shape = (h, w, 3)
             
-            # Convert to PIL Image
-            output_image = Image.fromarray(plot_data)
+            # Create PIL Image from buffer
+            output_image = Image.fromarray(buf)
             
-            # Resize if needed
-            output_image = output_image.resize(original_size, Image.Resampling.LANCZOS)
+            # Resize to original size if needed
+            if output_image.size != original_size:
+                output_image = output_image.resize(original_size, Image.Resampling.LANCZOS)
             
             # Save to final buffer
             final_buf = io.BytesIO()
@@ -294,13 +293,17 @@ def process_image_detection(image, target_label, surprise_rating):
             plt.close(fig)
             
             return final_buf
-            
+
         except Exception as e:
-            print(f"Save error details: {str(e)}")  # Debug print
+            print(f"Save error details: {str(e)}")
             print(f"Figure type: {type(fig)}")
             print(f"Canvas type: {type(fig.canvas)}")
             raise
 
+    except Exception as e:
+        print(f"Process image detection error: {str(e)}")
+        print(f"Error occurred at line {e.__traceback__.tb_lineno}")
+        raise
 
 def process_and_analyze(image):
     if image is None: