Update app.py

app.py

@@ -60,13 +60,20 @@ async def predict_ndvi_api(file: UploadFile = File(...)):
 async def predict_yolo_api(file: UploadFile = File(...)):
     """Predict YOLO results from 4-channel TIFF image"""
     try:
-        # Save uploaded file temporarily
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.tif') as tmp_file:
+        # Save uploaded file temporarily with proper extension
+        file_extension = '.tiff' if file.filename.lower().endswith(('.tif', '.tiff')) else '.tiff'
+        
+        with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp_file:
             contents = await file.read()
             tmp_file.write(contents)
+            tmp_file.flush()  # Ensure data is written
             tmp_file_path = tmp_file.name
         
         try:
+            # Verify the file was written correctly
+            if not os.path.exists(tmp_file_path) or os.path.getsize(tmp_file_path) == 0:
+                raise ValueError("Failed to create temporary file")
+            
             # Predict using YOLO model
             results = predict_yolo(yolo_model, tmp_file_path)
             
@@ -79,23 +86,24 @@ async def predict_yolo_api(file: UploadFile = File(...)):
                 },
                 "classes": results.boxes.cls.tolist() if results.boxes is not None else None,
                 "names": results.names,
-                "growth_stages": getattr(results, 'growth_stages', None),
                 "orig_shape": results.orig_shape,
                 "speed": results.speed
             }
             
-            # Handle growth stages if present
-            if hasattr(results, 'boxes') and hasattr(results.boxes, 'data'):
-                # Extract growth stages from the results if available
-                if len(results.boxes.data[0]) > 6:  # Assuming growth stages are in the data
-                    growth_stages = results.boxes.data[:, 6].tolist()
-                    results_dict["growth_stages"] = growth_stages
+            # Handle growth stages if present in the results
+            if hasattr(results, 'boxes') and results.boxes is not None:
+                if hasattr(results.boxes, 'data') and len(results.boxes.data) > 0:
+                    # Check if there are additional columns for growth stages
+                    if results.boxes.data.shape[1] > 6:
+                        growth_stages = results.boxes.data[:, 6:].tolist()
+                        results_dict["growth_stages"] = growth_stages
             
             return JSONResponse(content=results_dict)
             
         finally:
             # Clean up temporary file
-            os.unlink(tmp_file_path)
+            if os.path.exists(tmp_file_path):
+                os.unlink(tmp_file_path)
             
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e)})
@@ -104,13 +112,20 @@ async def predict_yolo_api(file: UploadFile = File(...)):
 async def predict_pipeline_api(file: UploadFile = File(...)):
     """Full pipeline: RGB -> NDVI -> 4-channel -> YOLO prediction"""
     try:
-        # Save uploaded file temporarily
-        with tempfile.NamedTemporaryFile(delete=False, suffix='.tif') as tmp_file:
+        # Save uploaded file temporarily with proper extension
+        file_extension = '.tiff' if file.filename.lower().endswith(('.tif', '.tiff')) else '.jpg'
+        
+        with tempfile.NamedTemporaryFile(delete=False, suffix=file_extension) as tmp_file:
             contents = await file.read()
             tmp_file.write(contents)
+            tmp_file.flush()  # Ensure data is written
             tmp_file_path = tmp_file.name
         
         try:
+            # Verify the file was written correctly
+            if not os.path.exists(tmp_file_path) or os.path.getsize(tmp_file_path) == 0:
+                raise ValueError("Failed to create temporary file")
+            
             # Run the full pipeline
             results = predict_pipeline(ndvi_model, yolo_model, tmp_file_path)
             
@@ -123,22 +138,24 @@ async def predict_pipeline_api(file: UploadFile = File(...)):
                 },
                 "classes": results.boxes.cls.tolist() if results.boxes is not None else None,
                 "names": results.names,
-                "growth_stages": getattr(results, 'growth_stages', None),
                 "orig_shape": results.orig_shape,
                 "speed": results.speed
             }
             
-            # Handle growth stages if present
-            if hasattr(results, 'boxes') and hasattr(results.boxes, 'data'):
-                if len(results.boxes.data[0]) > 6:
-                    growth_stages = results.boxes.data[:, 6].tolist()
-                    results_dict["growth_stages"] = growth_stages
+            # Handle growth stages if present in the results
+            if hasattr(results, 'boxes') and results.boxes is not None:
+                if hasattr(results.boxes, 'data') and len(results.boxes.data) > 0:
+                    # Check if there are additional columns for growth stages
+                    if results.boxes.data.shape[1] > 6:
+                        growth_stages = results.boxes.data[:, 6:].tolist()
+                        results_dict["growth_stages"] = growth_stages
             
             return JSONResponse(content=results_dict)
             
         finally:
             # Clean up temporary file
-            os.unlink(tmp_file_path)
+            if os.path.exists(tmp_file_path):
+                os.unlink(tmp_file_path)
             
     except Exception as e:
         return JSONResponse(status_code=500, content={"error": str(e)})