CLIPSeg

Sleeping

App Files Files Community

sigyllly commited on Dec 24, 2024

Commit

63596f0

verified ·

1 Parent(s): 90674f3

Update app.py

Browse files

Files changed (1) hide show

app.py +97 -139

app.py CHANGED Viewed

@@ -1,146 +1,104 @@
-from flask import Flask, request, jsonify, render_template_string
-import subprocess
 app = Flask(__name__)
-# Route to the homepage with embedded HTML
 @app.route('/')
 def index():
-    html = """
-    <!DOCTYPE html>
-    <html lang="en">
-    <head>
-        <meta charset="UTF-8">
-        <meta name="viewport" content="width=device-width, initial-scale=1.0">
-        <title>Administrator Terminal</title>
-        <style>
-            body {
-                font-family: Arial, sans-serif;
-                margin: 0;
-                padding: 0;
-                background-color: #f4f4f4;
-            }
-            .container {
-                width: 80%;
-                margin: 50px auto;
-                padding: 20px;
-                background-color: white;
-                box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
-                border-radius: 10px;
-            }
-            #terminal {
-                width: 100%;
-                height: 300px;
-                background-color: #000;
-                color: #0f0;
-                font-family: "Courier New", Courier, monospace;
-                padding: 10px;
-                overflow-y: scroll;
-                white-space: pre-wrap;
-                border-radius: 5px;
-                margin-bottom: 10px;
-            }
-            #command {
-                width: 100%;
-                padding: 10px;
-                font-size: 1em;
-                margin-bottom: 10px;
-                border: 1px solid #ccc;
-                border-radius: 5px;
-            }
-            button {
-                padding: 10px 20px;
-                font-size: 1em;
-                border: none;
-                border-radius: 5px;
-                background-color: #28a745;
-                color: white;
-                cursor: pointer;
-            }
-            button:hover {
-                background-color: #218838;
-            }
-        </style>
-    </head>
-    <body>
-        <div class="container">
-            <h2>Administrator Terminal</h2>
-            <div id="terminal"></div>
-            <input type="text" id="command" placeholder="Enter command..." autofocus>
-            <button onclick="sendCommand()">Execute</button>
-        </div>
-        <script>
-            // Function to send command to the Flask server
-            function sendCommand() {
-                const command = document.getElementById('command').value;
-                if (command.trim() === '') {
-                    alert('Please enter a command');
-                    return;
-                }
-                fetch('/execute', {
-                    method: 'POST',
-                    headers: {
-                        'Content-Type': 'application/x-www-form-urlencoded'
-                    },
-                    body: new URLSearchParams({
-                        'command': command
-                    })
-                })
-                .then(response => response.json())
-                .then(data => {
-                    const terminal = document.getElementById('terminal');
-                    if (data.stdout) {
-                        terminal.innerHTML += '> ' + command + '\\n' + data.stdout + '\\n';
-                    }
-                    if (data.stderr) {
-                        terminal.innerHTML += '> ' + command + '\\n' + data.stderr + '\\n';
-                    }
-                    if (data.error) {
-                        terminal.innerHTML += '> ' + command + '\\n' + data.error + '\\n';
-                    }
-                    document.getElementById('command').value = '';
-                    terminal.scrollTop = terminal.scrollHeight;  // Scroll to the bottom
-                })
-                .catch(error => {
-                    console.error('Error:', error);
-                });
-            }
-            // Allow pressing Enter to send the command
-            document.getElementById('command').addEventListener('keydown', function (e) {
-                if (e.key === 'Enter') {
-                    sendCommand();
-                }
-            });
-        </script>
-    </body>
-    </html>
-    """
-    return render_template_string(html)
-# Route to execute a command
-@app.route('/execute', methods=['POST'])
-def execute():
-    try:
-        # Get the command from the request
-        command = request.form['command']
-        # Execute the command and capture the output
-        result = subprocess.run(command, shell=True, capture_output=True, text=True)
-        # Return the output (stdout and stderr)
-        return jsonify({
-            'stdout': result.stdout,
-            'stderr': result.stderr
-        })
-    except Exception as e:
-        return jsonify({
-            'error': str(e)
-        })
 if __name__ == '__main__':
-     app.run(host='0.0.0.0', port=7860, debug=True)

+from flask import Flask, request, jsonify, render_template
+from PIL import Image
+import base64
+from io import BytesIO
+from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import cv2
 app = Flask(__name__)
+processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
+model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
+def process_image(image, prompt, threshold, alpha_value, draw_rectangles):
+    inputs = processor(
+        text=prompt, images=image, padding="max_length", return_tensors="pt"
+    )
+    # predict
+    with torch.no_grad():
+        outputs = model(**inputs)
+        preds = outputs.logits
+    pred = torch.sigmoid(preds)
+    mat = pred.cpu().numpy()
+    mask = Image.fromarray(np.uint8(mat * 255), "L")
+    mask = mask.convert("RGB")
+    mask = mask.resize(image.size)
+    mask = np.array(mask)[:, :, 0]
+    # normalize the mask
+    mask_min = mask.min()
+    mask_max = mask.max()
+    mask = (mask - mask_min) / (mask_max - mask_min)
+    # threshold the mask
+    bmask = mask > threshold
+    # zero out values below the threshold
+    mask[mask < threshold] = 0
+    fig, ax = plt.subplots()
+    ax.imshow(image)
+    ax.imshow(mask, alpha=alpha_value, cmap="jet")
+    if draw_rectangles:
+        contours, hierarchy = cv2.findContours(
+            bmask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+        )
+        for contour in contours:
+            x, y, w, h = cv2.boundingRect(contour)
+            rect = plt.Rectangle(
+                (x, y), w, h, fill=False, edgecolor="yellow", linewidth=2
+            )
+            ax.add_patch(rect)
+    ax.axis("off")
+    plt.tight_layout()
+    bmask = Image.fromarray(bmask.astype(np.uint8) * 255, "L")
+    output_image = Image.new("RGBA", image.size, (0, 0, 0, 0))
+    output_image.paste(image, mask=bmask)
+    # Convert mask to base64
+    buffered_mask = BytesIO()
+    bmask.save(buffered_mask, format="PNG")
+    result_mask = base64.b64encode(buffered_mask.getvalue()).decode('utf-8')
+    # Convert output image to base64
+    buffered_output = BytesIO()
+    output_image.save(buffered_output, format="PNG")
+    result_output = base64.b64encode(buffered_output.getvalue()).decode('utf-8')
+    return fig, result_mask, result_output
+    # Existing process_image function, copy it here
+    # ...
 @app.route('/')
 def index():
+    return render_template('index.html')
+@app.route('/api/mask_image', methods=['POST'])
+def mask_image_api():
+    data = request.get_json()
+    base64_image = data.get('base64_image', '')
+    prompt = data.get('prompt', '')
+    threshold = data.get('threshold', 0.4)
+    alpha_value = data.get('alpha_value', 0.5)
+    draw_rectangles = data.get('draw_rectangles', False)
+    # Decode base64 image
+    image_data = base64.b64decode(base64_image.split(',')[1])
+    image = Image.open(BytesIO(image_data))
+    # Process the image
+    _, result_mask, result_output = process_image(image, prompt, threshold, alpha_value, draw_rectangles)
+    return jsonify({'result_mask': result_mask, 'result_output': result_output})
 if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=7860, debug=True)