add: custom error hadling

app.py

@@ -22,4 +22,4 @@ app.add_middleware(
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
-)
+)

src/api/image_prep_api.py

@@ -49,173 +49,244 @@ def replicate_enhancer(input):
 
 @preprocessing_router.post("/rem_bg")
 async def remove_background(image: UploadFile = File(...)):
-    image_bytes = await image.read()
-    image = Image.open(BytesIO(image_bytes)).convert("RGB")
-    act_img_base_64 = BytesIO()
-    image.save(act_img_base_64, format="WEBP")
-    image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
+    start_time = time.time()
 
-    image_data_uri = f"data:image/WEBP;base64,{image_bytes_}"
+    try:
+        image_bytes = await image.read()
+        image = Image.open(BytesIO(image_bytes)).convert("RGB")
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": f"Error reading image: {str(e)}", "code": 500})
 
-    input = {
-        "image": image_data_uri,
-    }
-    output = replicate_bg(input)
+    try:
+        act_img_base_64 = BytesIO()
+        image.save(act_img_base_64, format="WEBP")
+        image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
+        image_data_uri = f"data:image/WEBP;base64,{image_bytes_}"
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error converting image to base64: {str(e)}", "code": 500})
 
-    response = requests.get(output)
-    base_64 = base64.b64encode(response.content).decode('utf-8')
-    base64_prefix = "data:image/WEBP;base64,"
+    try:
+        output = replicate_bg({"image": image_data_uri})
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error running background removal: {str(e)}", "code": 500})
 
     try:
+        response = requests.get(output)
+        base_64 = base64.b64encode(response.content).decode('utf-8')
+        base64_prefix = "data:image/WEBP;base64,"
+
+        total_inference_time = round((time.time() - start_time), 2)
+
         response = {
             "output": f"{base64_prefix}{base_64}",
-            'code': 200
+            "inference_time": total_inference_time,
+            "code": 200
         }
-
         return JSONResponse(content=response, status_code=200)
 
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to process image: {e}")
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error processing response: {str(e)}", "code": 500})
 
 
 @preprocessing_router.post("/upscale_image")
 async def upscale_image(image: UploadFile = File(...), scale: int = 1):
-    image_bytes = await image.read()
-    image = Image.open(BytesIO(image_bytes)).convert("RGBA")
-    act_img_base_64 = BytesIO()
-    image.save(act_img_base_64, format="PNG")
-    image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
-
-    image_data_uri = f"data:image/png;base64,{image_bytes_}"
+    start_time = time.time()
 
-    input = {
-        "image": image_data_uri,
-        "scale": scale,
-        "face_enhance": False
+    try:
+        image_bytes = await image.read()
+        image = Image.open(BytesIO(image_bytes)).convert("RGBA")
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": f"Error reading image: {str(e)}", "code": 500})
 
-    }
-    output = replicate_enhancer(input)
+    try:
+        act_img_base_64 = BytesIO()
+        image.save(act_img_base_64, format="PNG")
+        image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
+        image_data_uri = f"data:image/png;base64,{image_bytes_}"
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error converting image to base64: {str(e)}", "code": 500})
 
-    response = requests.get(output)
-    base_64 = base64.b64encode(response.content).decode('utf-8')
-    base64_prefix = image_data_uri.split(",")[0] + ","
+    try:
+        input = {
+            "image": image_data_uri,
+            "scale": scale,
+            "face_enhance": False
+        }
+        output = replicate_enhancer(input)
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error running image enhancement: {str(e)}", "code": 500})
 
     try:
+        response = requests.get(output)
+        base_64 = base64.b64encode(response.content).decode('utf-8')
+        base64_prefix = image_data_uri.split(",")[0] + ","
+
+        total_inference_time = round((time.time() - start_time), 2)
+
         response = {
             "output": f"{base64_prefix}{base_64}",
-            'code': 200
+            "inference_time": total_inference_time,
+            "code": 200
         }
-
         return JSONResponse(content=response, status_code=200)
 
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to process image: {e}")
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error processing response: {str(e)}", "code": 500})
 
 
 @preprocessing_router.post("/crop_transparent")
 async def crop_transparent(image: UploadFile):
-    if not image.content_type == "image/png":
-        raise HTTPException(status_code=400, detail="Only PNG files are supported")
+    start_time = time.time()
 
     try:
-        contents = await image.read()
+        if not image.content_type == "image/png":
+            return JSONResponse(status_code=400,
+                                content={"error": "Only PNG files are supported", "code": 400})
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error checking file type: {str(e)}", "code": 500})
 
+    try:
+        contents = await image.read()
         cropped_image_bytes, metadata = crop_transparent_image(contents)
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error cropping image: {str(e)}", "code": 500})
 
+    try:
         base64_image = base64.b64encode(cropped_image_bytes).decode('utf-8')
         base64_prefix = "data:image/png;base64,"
 
-        return {
+        total_inference_time = round((time.time() - start_time), 2)
+
+        return JSONResponse(content={
             "status": "success",
+            "code": 200,
             "data": {
                 "image": f"{base64_prefix}{base64_image}",
-                "metadata": metadata
+                "metadata": metadata,
+                "inference_time": total_inference_time
             }
-        }
-
-    except ValueError as e:
-        raise HTTPException(status_code=400, detail=str(e))
+        }, status_code=200)
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error processing response: {str(e)}", "code": 500})
 
 
 @preprocessing_router.post("/background_replace")
 async def bg_replace(image: UploadFile = File(...), bg_image: UploadFile = File(...)):
-    image_bytes = await image.read()
-    bg_image = await bg_image.read()
+    start_time = time.time()
 
-    image, bg_image = Image.open(BytesIO(image_bytes)).convert("RGBA"), Image.open(BytesIO(bg_image)).convert("RGB")
-    width, height = bg_image.size
-    background = Image.fromarray(np.array(bg_image)).resize((width, height))
-    orig_img = Image.fromarray(np.array(image)).resize((width, height))
-    background.paste(orig_img, (0, 0), mask=orig_img)
-    act_img_base_64 = BytesIO()
-    background.save(act_img_base_64, format="WEBP")
-    image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
+    try:
+        image_bytes = await image.read()
+        bg_bytes = await bg_image.read()
+        image = Image.open(BytesIO(image_bytes)).convert("RGBA")
+        bg_image = Image.open(BytesIO(bg_bytes)).convert("RGB")
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error reading images: {str(e)}", "code": 500})
 
-    image_data_uri = f"data:image/webp;base64,{image_bytes_}"
+    try:
+        width, height = bg_image.size
+        background = Image.fromarray(np.array(bg_image)).resize((width, height))
+        orig_img = Image.fromarray(np.array(image)).resize((width, height))
+        background.paste(orig_img, (0, 0), mask=orig_img)
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error processing images: {str(e)}", "code": 500})
 
     try:
-        response = {
-            "output": f"{image_data_uri}",
-            'code': 200
-        }
+        act_img_base_64 = BytesIO()
+        background.save(act_img_base_64, format="WEBP")
+        image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
+        image_data_uri = f"data:image/webp;base64,{image_bytes_}"
 
-        return JSONResponse(content=response, status_code=200)
+        total_inference_time = round((time.time() - start_time), 2)
 
+        return JSONResponse(content={
+            "output": image_data_uri,
+            "code": 200,
+            "inference_time": total_inference_time
+        }, status_code=200)
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to process image: {e}")
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error creating response: {str(e)}", "code": 500})
 
 
 @preprocessing_router.post("/rem_bg_color_extraction")
-async def remove_background_color_extraction(image: UploadFile = File(...), hex_color: str = "#FFFFFF",
+async def remove_background_color_extraction(image: UploadFile = File(...),
+                                             hex_color: str = "#FFFFFF",
                                              threshold: int = 30):
-    image_bytes = await image.read()
-    image = Image.open(BytesIO(image_bytes)).convert("RGBA")
-    image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    start_time = time.time()
 
-    result = color_extraction_rmbg.extract_color(image, hex_color, threshold)
-    result = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_RGB2BGRA)).convert("RGBA")
-    act_img_base_64 = BytesIO()
-    result.save(act_img_base_64, format="PNG")
-    image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
+    try:
+        image_bytes = await image.read()
+        image = Image.open(BytesIO(image_bytes)).convert("RGBA")
+        image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error reading image: {str(e)}", "code": 500})
 
-    image_data_uri = f"data:image/png;base64,{image_bytes_}"
+    try:
+        result = color_extraction_rmbg.extract_color(image, hex_color, threshold)
+        result = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_RGB2BGRA)).convert("RGBA")
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error extracting colors: {str(e)}", "code": 500})
 
     try:
-        response = {
-            "output": f"{image_data_uri}",
-            'code': 200
-        }
+        act_img_base_64 = BytesIO()
+        result.save(act_img_base_64, format="PNG")
+        image_bytes_ = base64.b64encode(act_img_base_64.getvalue()).decode("utf-8")
+        image_data_uri = f"data:image/png;base64,{image_bytes_}"
 
-        return JSONResponse(content=response, status_code=200)
+        total_inference_time = round((time.time() - start_time), 2)
 
+        return JSONResponse(content={
+            "output": image_data_uri,
+            "code": 200,
+            "inference_time": total_inference_time
+        }, status_code=200)
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Failed to process image: {e}")
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error creating response: {str(e)}", "code": 500})
 
 
 @preprocessing_router.post("/title_description_generator")
 async def product_title_description_generator(image: UploadFile = File(...)):
     start_time = time.time()
-    image_bytes = await image.read()
-    image = Image.open(BytesIO(image_bytes)).convert("RGB")
+
     try:
-        result = product_listing_obj.gen_title_desc(image=image)
-        print(result)
+        image_bytes = await image.read()
+        image = Image.open(BytesIO(image_bytes)).convert("RGB")
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error reading image: {str(e)}", "code": 500})
 
+    try:
+        result = product_listing_obj.gen_title_desc(image=image)
         title = result.split("Title:")[1].split("Description:")[0]
         description = result.split("Description:")[1]
-        inference_time = time.time() - start_time
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                            content={"error": "Please make sure the image is clear and necklaces are visible",
+                                     "code": 500})
+
+    try:
+        total_inference_time = round((time.time() - start_time), 2)
 
-        result = {
+        return JSONResponse(content={
             "code": 200,
             "title": title,
             "description": description,
-            "inference_time": inference_time
-        }
-
-        return JSONResponse(content=result, status_code=200)
-
+            "inference_time": total_inference_time
+        }, status_code=200)
     except Exception as e:
-        raise HTTPException(status_code=500,
-                            detail=f"Please make sure the image is clear and necklaces are visible{str(e)}")
+        return JSONResponse(status_code=500,
+                            content={"error": f"Error creating response: {str(e)}", "code": 500})

src/api/image_regeneration_api.py

@@ -31,30 +31,26 @@ async def image_re_gen(
         negative_prompt: str = Form(""),
         image: UploadFile = File(...),
         mask_image: Optional[UploadFile] = File(default=None),
-
         reference_image_c1: Optional[UploadFile] = File(default=None),
         reference_image_c1_type: Optional[str] = Form(default=""),
         reference_image_c1_weight: Optional[float] = Form(default=0.0),
         reference_image_c1_stop: Optional[float] = Form(default=0.0),
-
         reference_image_c2: Optional[UploadFile] = File(default=None),
         reference_image_c2_type: Optional[str] = Form(default=""),
         reference_image_c2_weight: Optional[float] = Form(default=0.0),
         reference_image_c2_stop: Optional[float] = Form(default=0.0),
-
         reference_image_c3: Optional[UploadFile] = File(default=None),
         reference_image_c3_type: Optional[str] = Form(default=""),
         reference_image_c3_weight: Optional[float] = Form(default=0.0),
         reference_image_c3_stop: Optional[float] = Form(default=0.0),
-
         reference_image_c4: Optional[UploadFile] = File(default=None),
         reference_image_c4_type: Optional[str] = Form(default=""),
         reference_image_c4_weight: Optional[float] = Form(default=0.0),
         reference_image_c4_stop: Optional[float] = Form(default=0.0),
 ):
     start_time = time.time()
-    try:
 
+    try:
         async def process_reference_image(reference_image: Optional[UploadFile]) -> Optional[str]:
             if reference_image is not None:
                 reference_image_bytes = await reference_image.read()
@@ -64,22 +60,32 @@ async def image_re_gen(
                 reference_image_b64 = base64.b64encode(ref_img_base64.getvalue()).decode("utf-8")
                 return f"data:image/WEBP;base64,{reference_image_b64}"
             return None
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                          content={"error": f"Error processing reference image: {str(e)}", "code": 500})
 
-        # Process main image
+    try:
         image_bytes = await image.read()
         image = Image.open(BytesIO(image_bytes)).convert("RGB")
         img_base64 = BytesIO()
         image.save(img_base64, format="WEBP")
         image_data_uri = f"data:image/WEBP;base64,{base64.b64encode(img_base64.getvalue()).decode('utf-8')}"
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                          content={"error": f"Error processing main image: {str(e)}", "code": 500})
 
+    try:
         reference_images = {
             'c1': await process_reference_image(reference_image_c1),
             'c2': await process_reference_image(reference_image_c2),
             'c3': await process_reference_image(reference_image_c3),
             'c4': await process_reference_image(reference_image_c4)
         }
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                          content={"error": f"Error processing reference images: {str(e)}", "code": 500})
 
-        # Initialize input data with required fields
+    try:
         input_data = {
             "prompt": prompt,
             "inpaint_input_image": image_data_uri,
@@ -100,7 +106,11 @@ async def image_re_gen(
             mask_image.save(mask_base64, format="WEBP")
             mask_image_data_uri = f"data:image/WEBP;base64,{base64.b64encode(mask_base64.getvalue()).decode('utf-8')}"
             input_data["inpaint_input_mask"] = mask_image_data_uri
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                          content={"error": f"Error preparing input data: {str(e)}", "code": 500})
 
+    try:
         for i in range(1, 5):
             c = f'c{i}'
             if reference_images[c] is not None:
@@ -117,19 +127,27 @@ async def image_re_gen(
                 stop_value = locals()[f'reference_image_{c}_stop']
                 if stop_value != 0.0 or stop_value != 0:
                     input_data[f"cn_stop{i}"] = stop_value
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                          content={"error": f"Error processing reference image parameters: {str(e)}", "code": 500})
 
+    try:
         output = image_regeneration_replicate(input_data)
         response = requests.get(output[0])
         output_base64 = base64.b64encode(response.content).decode('utf-8')
         base64_prefix = image_data_uri.split(",")[0] + ","
+    except Exception as e:
+        return JSONResponse(status_code=500,
+                          content={"error": f"Error generating image: {str(e)}", "code": 500})
+
+    try:
         inference_time = time.time() - start_time
         response = {
             "output": f"{base64_prefix}{output_base64}",
             "inference_time": inference_time,
             "code": 200,
         }
-
         return JSONResponse(content=response, status_code=200)
     except Exception as e:
-        return JSONResponse(content={"code": 500, "error": "Please try again later there might be some issue"},
-                            status_code=500)
+        return JSONResponse(status_code=500,
+                          content={"error": f"Error creating response: {str(e)}", "code": 500})