Update app.py

app.py

@@ -6,7 +6,6 @@ from tensorflow.keras.models import load_model
 from PIL import Image
 import io
 import os
-import cv2
 
 app = FastAPI(title="GAN Image Generator API")
 
@@ -26,22 +25,21 @@ generator = load_model(model_path)
 def preprocess_sketch(image_bytes):
     """Process uploaded sketch image for the GAN"""
     try:
-        # Convert bytes to PIL Image
+        # Convert bytes to PIL Image and ensure it's RGB
         img = Image.open(io.BytesIO(image_bytes)).convert('RGB')
-        img = np.array(img)
-        
-        # Convert to grayscale
-        gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-        
-        # Resize to model's expected input size (adjust dimensions as needed)
-        resized = cv2.resize(gray, (256, 256))
-        
-        # Normalize pixel values to [-1, 1] range
-        normalized = (resized.astype(np.float32) / 127.5) - 1.0
-        
-        # Add batch and channel dimensions
-        processed = np.expand_dims(normalized, axis=[0, -1])
         
+        # Resize to model's expected input size
+        img = img.resize((256, 256))
+
+        # Convert to numpy array and normalize
+        img_array = np.array(img).astype(np.float32) / 255.0
+
+        # Ensure the array has 3 channels
+        if len(img_array.shape) == 2:
+            img_array = np.stack((img_array,) * 3, axis=-1)
+
+        # Add batch dimension (1, 256, 256, 3)
+        processed = np.expand_dims(img_array, axis=0)
         return processed
     except Exception as e:
         raise ValueError(f"Image processing failed: {str(e)}")
@@ -51,21 +49,23 @@ async def generate_from_sketch(file: UploadFile = File(...)):
     try:
         # Read uploaded file
         contents = await file.read()
-        
+
         # Process sketch
         processed_sketch = preprocess_sketch(contents)
-        
-        # Generate image (modify this to use your actual GAN prediction)
+
+        # Generate image using GAN
         generated = generator.predict(processed_sketch)
-        generated = (generated.squeeze() * 255).astype(np.uint8)
-        
+        generated = np.clip(generated[0], 0, 1) * 255
+        generated = generated.astype(np.uint8)
+
         # Convert to bytes
         img = Image.fromarray(generated)
         img_byte_arr = io.BytesIO()
         img.save(img_byte_arr, format='PNG')
         img_byte_arr.seek(0)
-        
+
         return StreamingResponse(img_byte_arr, media_type="image/png")
-        
+
     except Exception as e:
-        raise HTTPException(status_code=400, detail=str(e))
+        raise HTTPException(status_code=400, detail=str(e))
+