doublelotus/Removal

1from flask import Flask, request, send_file, Response, jsonify2from flask_cors import CORS3import numpy as np4import io5import torch6import cv27from segment_anything import sam_model_registry, SamAutomaticMaskGenerator8from PIL import Image9import zipfile10 11app = Flask(__name__)12CORS(app)13 14cudaOrNah = "cuda" if torch.cuda.is_available() else "cpu"15print(cudaOrNah)16 17# Global model setup 18# running out of memory adjusted19# checkpoint = "sam_vit_h_4b8939.pth"20# model_type = "vit_h"21checkpoint = "sam_vit_l_0b3195.pth"22model_type = "vit_l"23sam = sam_model_registry[model_type](checkpoint=checkpoint)24sam.to(device=cudaOrNah)25mask_generator = SamAutomaticMaskGenerator(26    model=sam,27    min_mask_region_area=0.0015  # Adjust this value as needed28)29print('Setup SAM model')30 31@app.route('/')32def hello():33    return {"hei": "Shredded to peices"}34 35@app.route('/health', methods=['GET'])36def health_check():37    # Simple health check endpoint38    return jsonify({"status": "ok"}), 20039 40@app.route('/get-masks', methods=['POST'])41def get_masks():42    try:43        print('received image from frontend')44        # Get the image file from the request45        if 'image' not in request.files:46            return jsonify({"error": "No image file provided"}), 40047        48        image_file = request.files['image']49        if image_file.filename == '':50            return jsonify({"error": "No image file provided"}), 40051 52        raw_image = Image.open(image_file).convert("RGB")53        # Convert the PIL Image to a NumPy array54        image_array = np.array(raw_image)55        # Since OpenCV expects BGR, convert RGB to BGR56        image = image_array[:, :, ::-1]57 58        if image is None:59            raise ValueError("Image not found or unable to read.")60        61        if cudaOrNah == "cuda":62            torch.cuda.empty_cache()63        64        masks = mask_generator.generate(image)65 66        if cudaOrNah == "cuda":67            torch.cuda.empty_cache()68 69        masks = sorted(masks, key=(lambda x: x['area']), reverse=True)70 71        def is_background(segmentation):72            val = (segmentation[10, 10] or segmentation[-10, 10] or73                   segmentation[10, -10] or segmentation[-10, -10])74            return val75 76        masks = [mask for mask in masks if not is_background(mask['segmentation'])]77 78        for i in range(0, len(masks) - 1)[::-1]:79            large_mask = masks[i]['segmentation']80            for j in range(i+1, len(masks)):81                not_small_mask = np.logical_not(masks[j]['segmentation'])82                masks[i]['segmentation'] = np.logical_and(large_mask, not_small_mask)83                masks[i]['area'] = masks[i]['segmentation'].sum()84                large_mask = masks[i]['segmentation']85 86        def sum_under_threshold(segmentation, threshold):87            return segmentation.sum() / segmentation.size < 0.001588 89        masks = [mask for mask in masks if not sum_under_threshold(mask['segmentation'], 100)]90        masks = sorted(masks, key=(lambda x: x['area']), reverse=True)91 92        # Create a zip file in memory93        zip_buffer = io.BytesIO()94        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:95            for idx, mask in enumerate(masks):96                alpha = mask['segmentation'].astype('uint8') * 25597                mask_image = Image.fromarray(alpha)98                mask_io = io.BytesIO()99                mask_image.save(mask_io, format="PNG")100                mask_io.seek(0)101                zip_file.writestr(f'mask_{idx+1}.png', mask_io.read())102 103        zip_buffer.seek(0)104        105        return send_file(zip_buffer, mimetype='application/zip', as_attachment=True, download_name='masks.zip')106    except Exception as e:107        # Log the error message if needed108        print(f"Error processing the image: {e}")109        # Return a JSON response with the error message and a 400 Bad Request status110        return jsonify({"error": "Error processing the image", "details": str(e)}), 400111 112if __name__ == '__main__':113    app.run(debug=True)114