MahmoudNabilMohamed commited on
Commit
87be849
ยท
verified ยท
1 Parent(s): 41ff11f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +217 -0
app.py ADDED
@@ -0,0 +1,217 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import numpy as np
3
+ import requests
4
+ from flask import Flask, request, jsonify, send_file
5
+ from werkzeug.utils import secure_filename
6
+ import onnxruntime
7
+ from PIL import Image
8
+ import io
9
+ import base64
10
+ import logging
11
+
12
+ # ุฅุนุฏุงุฏ ุงู„ุชุณุฌูŠู„
13
+ logging.basicConfig(level=logging.INFO)
14
+ logger = logging.getLogger(__name__)
15
+
16
+ app = Flask(__name__)
17
+ app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024 # 16MB max file size
18
+
19
+ # ุชุญุฏูŠุซ ุงู„ู…ุณุงุฑุงุช ู„ู„ุนู…ู„ ููŠ Codespaces
20
+ MODEL_URL = "https://huggingface.co/skytnt/anime-seg/resolve/main/isnetis.onnx"
21
+ MODEL_PATH = "isnetis.onnx"
22
+ UPLOAD_FOLDER = "uploads"
23
+ RESULT_FOLDER = "results"
24
+
25
+ # ุฅู†ุดุงุก ุงู„ู…ุฌู„ุฏุงุช
26
+ os.makedirs(UPLOAD_FOLDER, exist_ok=True)
27
+ os.makedirs(RESULT_FOLDER, exist_ok=True)
28
+
29
+ class BackgroundRemover:
30
+ def __init__(self):
31
+ self.session = None
32
+ self.input_name = None
33
+ self.output_name = None
34
+
35
+ def download_model(self):
36
+ if os.path.exists(MODEL_PATH):
37
+ logger.info("โœ… Model already exists, skipping download")
38
+ return True
39
+
40
+ try:
41
+ logger.info("๐Ÿ“ฅ Downloading model from HuggingFace...")
42
+ response = requests.get(MODEL_URL, stream=True)
43
+ response.raise_for_status()
44
+
45
+ total_size = int(response.headers.get('content-length', 0))
46
+ downloaded = 0
47
+
48
+ with open(MODEL_PATH, 'wb') as f:
49
+ for chunk in response.iter_content(chunk_size=8192):
50
+ if chunk:
51
+ f.write(chunk)
52
+ downloaded += len(chunk)
53
+ if total_size > 0:
54
+ progress = (downloaded / total_size) * 100
55
+ print(f"\r๐Ÿ“ฅ Downloading: {progress:.1f}%", end="", flush=True)
56
+
57
+ print() # ุณุทุฑ ุฌุฏูŠุฏ
58
+ logger.info("โœ… Model downloaded successfully")
59
+ return True
60
+ except Exception as e:
61
+ logger.error(f"โŒ Failed to download model: {e}")
62
+ return False
63
+
64
+ def load_model(self):
65
+ try:
66
+ logger.info("๐Ÿ”„ Loading ONNX model...")
67
+ providers = ['CPUExecutionProvider']
68
+
69
+ # ุงู„ุชุญู‚ู‚ ู…ู† ุฏุนู… GPU (ุงุฎุชูŠุงุฑูŠ)
70
+ try:
71
+ available_providers = onnxruntime.get_available_providers()
72
+ if 'CUDAExecutionProvider' in available_providers:
73
+ providers.insert(0, 'CUDAExecutionProvider')
74
+ logger.info("๐Ÿš€ CUDA provider available")
75
+ else:
76
+ logger.info("๐Ÿ’ป Using CPU provider")
77
+ except:
78
+ logger.info("๐Ÿ’ป Using CPU provider")
79
+
80
+ self.session = onnxruntime.InferenceSession(MODEL_PATH, providers=providers)
81
+ self.input_name = self.session.get_inputs()[0].name
82
+ self.output_name = self.session.get_outputs()[0].name
83
+
84
+ logger.info(f"โœ… Model loaded successfully")
85
+ logger.info(f"๐Ÿ“Š Input: {self.input_name}, Output: {self.output_name}")
86
+ return True
87
+ except Exception as e:
88
+ logger.error(f"โŒ Failed to load model: {e}")
89
+ return False
90
+
91
+ def preprocess_image(self, image):
92
+ original_size = image.size
93
+ image = image.convert('RGB')
94
+ image = image.resize((1024, 1024), Image.LANCZOS)
95
+ image_array = np.array(image).astype(np.float32)
96
+ image_array = image_array / 255.0
97
+ image_array = np.transpose(image_array, (2, 0, 1))
98
+ image_array = np.expand_dims(image_array, axis=0)
99
+ return image_array, original_size
100
+
101
+ def postprocess_mask(self, mask, original_size):
102
+ mask = mask.squeeze()
103
+ mask = (mask * 255).astype(np.uint8)
104
+ mask = Image.fromarray(mask, mode='L')
105
+ mask = mask.resize(original_size, Image.LANCZOS)
106
+ return mask
107
+
108
+ def remove_background(self, image_path):
109
+ try:
110
+ logger.info(f"๐Ÿ–ผ๏ธ Processing image: {image_path}")
111
+ image = Image.open(image_path)
112
+ preprocessed, original_size = self.preprocess_image(image)
113
+
114
+ logger.info("๐Ÿค– Running AI inference...")
115
+ mask = self.session.run([self.output_name], {self.input_name: preprocessed})[0]
116
+ mask = self.postprocess_mask(mask, original_size)
117
+
118
+ logger.info("โœ‚๏ธ Applying mask to remove background...")
119
+ image = image.convert('RGBA')
120
+ image.putalpha(mask)
121
+
122
+ logger.info("โœ… Background removed successfully")
123
+ return image
124
+ except Exception as e:
125
+ logger.error(f"โŒ Failed to remove background: {e}")
126
+ return None
127
+
128
+ # ุฅู†ุดุงุก ูƒุงุฆู† ุฅุฒุงู„ุฉ ุงู„ุฎู„ููŠุฉ
129
+ background_remover = BackgroundRemover()
130
+
131
+ @app.route('/', methods=['GET'])
132
+ def health_check():
133
+ """ูุญุต ุตุญุฉ ุงู„ุณูŠุฑูุฑ"""
134
+ return jsonify({
135
+ 'status': 'โœ… Server is running',
136
+ 'model_loaded': background_remover.session is not None,
137
+ 'endpoints': {
138
+ 'health_check': '/ (GET)',
139
+ 'remove_background': '/remove-background (POST)'
140
+ },
141
+ 'info': 'Background Remover API - Ready to use!'
142
+ })
143
+
144
+ @app.route('/remove-background', methods=['POST'])
145
+ def remove_background_endpoint():
146
+ """ุฅุฒุงู„ุฉ ุฎู„ููŠุฉ ุงู„ุตูˆุฑุฉ"""
147
+ try:
148
+ # ุงู„ุชุญู‚ู‚ ู…ู† ูˆุฌูˆุฏ ุงู„ู…ู„ู
149
+ if 'file' not in request.files:
150
+ return jsonify({'error': 'No file uploaded'}), 400
151
+
152
+ file = request.files['file']
153
+ if file.filename == '':
154
+ return jsonify({'error': 'No file selected'}), 400
155
+
156
+ # ุงู„ุชุญู‚ู‚ ู…ู† ู†ูˆุน ุงู„ู…ู„ู
157
+ if not file.content_type.startswith('image/'):
158
+ return jsonify({'error': 'File must be an image'}), 400
159
+
160
+ # ุญูุธ ุงู„ู…ู„ู
161
+ filename = secure_filename(file.filename)
162
+ filepath = os.path.join(UPLOAD_FOLDER, filename)
163
+ file.save(filepath)
164
+
165
+ logger.info(f"๐Ÿ“‚ File saved: {filepath}")
166
+
167
+ # ู…ุนุงู„ุฌุฉ ุงู„ุตูˆุฑุฉ
168
+ result = background_remover.remove_background(filepath)
169
+
170
+ if result is None:
171
+ os.remove(filepath) # ุญุฐู ุงู„ู…ู„ู ููŠ ุญุงู„ุฉ ุงู„ูุดู„
172
+ return jsonify({'error': 'Failed to process image'}), 500
173
+
174
+ # ุญูุธ ุงู„ู†ุชูŠุฌุฉ
175
+ result_filename = f"result_{filename.rsplit('.', 1)[0]}.png"
176
+ result_path = os.path.join(RESULT_FOLDER, result_filename)
177
+ result.save(result_path, 'PNG')
178
+
179
+ # ุญุฐู ุงู„ู…ู„ู ุงู„ุฃุตู„ูŠ
180
+ os.remove(filepath)
181
+
182
+ logger.info(f"โœ… Result saved: {result_path}")
183
+
184
+ return send_file(result_path, mimetype='image/png', as_attachment=False)
185
+
186
+ except Exception as e:
187
+ logger.error(f"โŒ Error in remove_background_endpoint: {e}")
188
+ return jsonify({'error': 'Internal server error'}), 500
189
+
190
+ if __name__ == '__main__':
191
+ print("=" * 60)
192
+ print("๐ŸŽฏ Background Remover API Server")
193
+ print("=" * 60)
194
+
195
+ logger.info("๐Ÿš€ Starting Background Remover Server...")
196
+
197
+ # ุชุญู…ูŠู„ ูˆุชุดุบูŠู„ ุงู„ู†ู…ูˆุฐุฌ
198
+ if not background_remover.download_model():
199
+ logger.error("โŒ Failed to download model. Exiting...")
200
+ exit(1)
201
+
202
+ if not background_remover.load_model():
203
+ logger.error("โŒ Failed to load model. Exiting...")
204
+ exit(1)
205
+
206
+ # ุฅุนุฏุงุฏ ุงู„ู…ู†ูุฐ
207
+ port = int(os.environ.get('PORT', 5000))
208
+
209
+ print(f"๐Ÿ“ก Running on port: {port}")
210
+ print(f"๐ŸŒ Local URL: http://localhost:{port}")
211
+ print(f"๐Ÿ” Health check: http://localhost:{port}/")
212
+ print(f"๐Ÿ“ฑ API endpoint: http://localhost:{port}/remove-background")
213
+ print("๐Ÿ’ก In Codespaces: Go to PORTS tab and make port 5000 public")
214
+ print("=" * 60)
215
+
216
+ # ุชุดุบูŠู„ ุงู„ุณูŠุฑูุฑ
217
+ app.run(host='0.0.0.0', port=port, debug=False, threaded=True)