Upload app.py with huggingface_hub

app.py

@@ -59,20 +59,54 @@ def initialize_classifier():
         logger.error(f"Failed to initialize classifier: {str(e)}")
         raise RuntimeError(f"Failed to initialize classifier: {str(e)}")
 
-# Validate image
-def validate_image(image):
-    if image is None:
-        return False, "No image provided"
-    if not isinstance(image, np.ndarray):
-        image = np.array(image)
-    if image.shape[0] < 10 or image.shape[1] < 10:
-        return False, "Image too small"
-    return True, image
+# Enhanced image validation and preprocessing
+def validate_and_preprocess_image(image):
+    try:
+        if image is None:
+            return False, "No image provided"
+        
+        # Convert PIL to numpy if needed
+        if hasattr(image, 'convert'):
+            image = image.convert('RGB')
+            image = np.array(image)
+        
+        if not isinstance(image, np.ndarray):
+            image = np.array(image)
+        
+        # Basic validation
+        if len(image.shape) < 2:
+            return False, "Invalid image format"
+        
+        if image.shape[0] < 10 or image.shape[1] < 10:
+            return False, "Image too small (minimum 10x10 pixels)"
+        
+        # Ensure 3-channel RGB
+        if len(image.shape) == 2:
+            image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+        elif len(image.shape) == 3 and image.shape[2] == 4:
+            image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
+        elif len(image.shape) == 3 and image.shape[2] == 3:
+            # Already RGB, but ensure proper format
+            pass
+        else:
+            return False, f"Unsupported image shape: {image.shape}"
+        
+        # Normalize to 0-255 range if needed
+        if image.dtype == np.float32 or image.dtype == np.float64:
+            if image.max() <= 1.0:
+                image = (image * 255).astype(np.uint8)
+        
+        logger.info(f"Preprocessed image shape: {image.shape}, dtype: {image.dtype}")
+        return True, image
+        
+    except Exception as e:
+        logger.error(f"Image preprocessing failed: {str(e)}")
+        return False, f"Image preprocessing failed: {str(e)}"
 
 # Core prediction function
 def predict_core(image, return_raw=False):
     try:
-        is_valid, result = validate_image(image)
+        is_valid, result = validate_and_preprocess_image(image)
         if not is_valid:
             logger.error(f"Invalid image: {result}")
             return {"error": result, "status": "failed"}
@@ -80,11 +114,7 @@ def predict_core(image, return_raw=False):
         image = result
         logger.info(f"Input image shape: {image.shape}, dtype: {image.dtype}")
         
-        # Ensure image is in RGB format
-        if len(image.shape) == 3 and image.shape[2] == 3:
-            # Convert RGB to BGR for OpenCV
-            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-        
+        # Use the classifier's predict method
         result = classifier.predict(image, return_probs=True)
         logger.info(f"Raw prediction result: {result}")
         
@@ -106,22 +136,19 @@ def predict_core(image, return_raw=False):
         logger.info(f"Prediction output: {output}")
         return output
     except Exception as e:
-        logger.error(f"Prediction failed: {str(e)}\n{traceback.format_exc()}")
-        return {"error": f"{str(e)}\n{traceback.format_exc()}", "status": "failed"}
+        error_msg = f"Prediction failed: {str(e)}"
+        logger.error(f"{error_msg}\n{traceback.format_exc()}")
+        return {"error": error_msg, "status": "failed"}
 
 # UI prediction function
 def predict(image):
     if image is None:
         return "Error: No image provided"
     
-    # Convert PIL Image to numpy array
-    if hasattr(image, 'convert'):
-        image = image.convert('RGB')
-        image = np.array(image)
-    
     result = predict_core(image)
     if result["status"] == "failed":
         return f"Error: {result['error']}"
+    
     return (
         f"**Classified as**: {result['prediction']}\n"
         f"**Confidence**: {result['confidence']:.4f}\n"
@@ -135,11 +162,6 @@ def predict_raw(image):
     if image is None:
         return json.dumps({"error": "No image provided", "status": "failed"}, indent=2)
     
-    # Convert PIL Image to numpy array
-    if hasattr(image, 'convert'):
-        image = image.convert('RGB')
-        image = np.array(image)
-    
     result = predict_core(image, return_raw=True)
     return json.dumps(result, indent=2)
 
@@ -156,21 +178,14 @@ app = FastAPI()
 @app.post("/api/predict/")
 async def predict_api(file: UploadFile = File(...)):
     try:
-        # Validate file type
         if not file.content_type.startswith('image/'):
             raise HTTPException(status_code=400, detail="File must be an image")
         
-        # Read file contents
         contents = await file.read()
-        
-        # Convert to PIL Image first
         image = Image.open(io.BytesIO(contents))
         image = image.convert('RGB')
         
-        # Convert to numpy array
-        image_array = np.array(image)
-        
-        result = predict_core(image_array)
+        result = predict_core(image)
         return JSONResponse(content=result)
     except Exception as e:
         logger.error(f"API prediction failed: {str(e)}")
@@ -182,21 +197,14 @@ async def predict_api(file: UploadFile = File(...)):
 @app.post("/api/predict_raw/")
 async def predict_raw_api(file: UploadFile = File(...)):
     try:
-        # Validate file type
         if not file.content_type.startswith('image/'):
             raise HTTPException(status_code=400, detail="File must be an image")
         
-        # Read file contents
         contents = await file.read()
-        
-        # Convert to PIL Image first
         image = Image.open(io.BytesIO(contents))
         image = image.convert('RGB')
         
-        # Convert to numpy array
-        image_array = np.array(image)
-        
-        result = predict_core(image_array, return_raw=True)
+        result = predict_core(image, return_raw=True)
         return JSONResponse(content=result)
     except Exception as e:
         logger.error(f"API raw prediction failed: {str(e)}")