Fix handling unsupported file types

app.py

@@ -277,36 +277,46 @@ def server(input, output, session: Session):
                 im = cv2.imread(file["datapath"])
 
                 # Convert BGR to RGB for display
-                im_rgb = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
-                pil_img = Image.fromarray(im_rgb)
-
-                # Convert to base64 for HTML display
-                buffered = BytesIO()
-                pil_img.save(buffered, format="PNG")
-                img_base64 = base64.b64encode(buffered.getvalue()).decode()
-
-                # Run prediction with original BGR image
-                prediction = predictor(im)
-                print(f"Analyzing image {idx + 1} of {len(files)}")
-                print(f"NMS threshold: {input.nms_threshold()}")
-                print(f'Number of instances: {len(prediction["instances"])}')
-                prediction = apply_nms(prediction, mask=True, cls_agnostic_nms=input.nms_threshold())
-                print(f'Number of instances after NMS: {len(prediction["instances"])}')
-
-                classes = prediction["instances"].pred_classes.tolist()
-
-                single_result = {
-                    "filename": file["name"],
-                    "image_base64": img_base64,
-                    "image": im,
-                    **prediction,
-                    "viable": classes.count(0),
-                    "non-viable": classes.count(1),
-                    "empty": classes.count(2),
-                    "total": len(classes),
-                    'NMS threshold': input.nms_threshold()
-                }
-                results.append(single_result)
+                try:
+                    im_rgb = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
+                    pil_img = Image.fromarray(im_rgb)
+
+                    # Convert to base64 for HTML display
+                    buffered = BytesIO()
+                    pil_img.save(buffered, format="PNG")
+                    img_base64 = base64.b64encode(buffered.getvalue()).decode()
+
+                    # Run prediction with original BGR image
+                    prediction = predictor(im)
+                    print(f"Analyzing image {idx + 1} of {len(files)}")
+                    print(f"NMS threshold: {input.nms_threshold()}")
+                    print(f'Number of instances: {len(prediction["instances"])}')
+                    prediction = apply_nms(prediction, mask=True, cls_agnostic_nms=input.nms_threshold())
+                    print(f'Number of instances after NMS: {len(prediction["instances"])}')
+
+                    classes = prediction["instances"].pred_classes.tolist()
+
+                    single_result = {
+                        "filename": file["name"],
+                        "image_base64": img_base64,
+                        "image": im,
+                        **prediction,
+                        "viable": classes.count(0),
+                        "non-viable": classes.count(1),
+                        "empty": classes.count(2),
+                        "total": len(classes),
+                        'NMS threshold': input.nms_threshold()
+                    }
+                    results.append(single_result)
+
+                except cv2.error as e:
+                    print(f"Error reading image {file['name']}: {e}")
+                    single_result = {
+                        "filename": file["name"]
+                    }
+                    results.append(single_result)
+
+
                 # print(f'Size of result: {sys.getsizeof(single_result)} bytes')
 
         # Update reactive value
@@ -378,7 +388,7 @@ def server(input, output, session: Session):
 
         # Create in-memory CSV file
         with tempfile.NamedTemporaryFile(delete=False, suffix=".csv") as tmp:
-            print(f'result tmp csv: {tmp.name}')
+            # print(f'result tmp csv: {tmp.name}')
             df.to_csv(tmp.name, index=False)
             return tmp.name
 
@@ -389,13 +399,16 @@ def server(input, output, session: Session):
         tmp_img_files = []
 
         with tempfile.TemporaryDirectory() as temp_dir:
-            print(os.listdir(os.path.dirname(temp_dir)))
+            # print(os.listdir(os.path.dirname(temp_dir)))
             for r in results:
                 # open your files here
                 named_file = os.path.join(temp_dir, r['filename'])
-                img = get_overlayed_image_from_single_result(r)
-                img.save(named_file)
-                tmp_img_files.append(named_file)
+                try:
+                    img = get_overlayed_image_from_single_result(r)
+                    img.save(named_file)
+                    tmp_img_files.append(named_file)
+                except KeyError as e:
+                    print(f"Error reading image {r['filename']}: {e}")
 
             with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as tmp: