Update app.py

app.py

@@ -1,51 +1,71 @@
-from fastapi import FastAPI, UploadFile, File
-from fastapi.responses import JSONResponse, StreamingResponse
-import numpy as np
-from tensorflow.keras.models import load_model
-from PIL import Image
-import io
-import os
-
-app = FastAPI(title="GAN Image Generator API")
-
-# Load model
-model_path = os.getenv("MODEL_PATH", "generator_final.h5")
-generator = load_model(model_path)
-
-def generate_image(noise_dim=100):
-    """Generate image from random noise"""
-    z = np.random.randn(1, noise_dim, 1, 1).astype(np.float32)
-    fake_image = generator.predict(z)
-    fake_image = (fake_image.squeeze() * 255).astype(np.uint8)
-    return fake_image[..., :3]
-
-@app.get("/")
-def read_root():
-    return {"message": "GAN Image Generator API"}
-
-@app.get("/generate-random")
-async def generate_random_image():
-    """Endpoint to generate random image"""
-    image_array = generate_image()
-    img = Image.fromarray(image_array)
-    
-    # Convert to bytes
-    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")
-
-@app.post("/generate-from-sketch")
-async def generate_from_sketch(file: UploadFile = File(...)):
-    """Endpoint to generate from sketch"""
-    # Process your sketch here (add your sketch processing logic)
-    # For now just returns a random image
-    image_array = generate_image()
-    img = Image.fromarray(image_array)
-    
-    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")
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse
+import numpy as np
+from tensorflow.keras.models import load_model
+from PIL import Image
+import io
+import os
+import cv2
+
+app = FastAPI(title="GAN Image Generator API")
+
+# Add CORS middleware to allow React frontend to connect
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins - adjust for production!
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Load model
+model_path = os.getenv("MODEL_PATH", "generator_final.h5")
+generator = load_model(model_path)
+
+def preprocess_sketch(image_bytes):
+    """Process uploaded sketch image for the GAN"""
+    try:
+        # Convert bytes to PIL Image
+        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])
+        
+        return processed
+    except Exception as e:
+        raise ValueError(f"Image processing failed: {str(e)}")
+
+@app.post("/generate-from-sketch")
+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)
+        generated = generator.predict(processed_sketch)
+        generated = (generated.squeeze() * 255).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))