Update app.py

app.py

@@ -14,7 +14,6 @@ from realesrgan.archs.srvgg_arch import SRVGGNetCompact
 # Global (CPU-only data; KHÔNG chạm CUDA ở đây)
 # --------------------
 last_file = None
-img_mode = "RGBA"
 
 DEVICE = "cpu"      # set trong gpu_startup()
 USE_HALF = False    # set trong gpu_startup()
@@ -49,28 +48,10 @@ def rnd_string(x):
     chars = "abcdefghijklmnopqrstuvwxyz_0123456789"
     return "".join(random.choice(chars) for _ in range(x))
 
-def has_transparency(img):
-    if img.info.get("transparency", None) is not None:
-        return True
-    if img.mode == "P":
-        transparent = img.info.get("transparency", -1)
-        for _, index in img.getcolors():
-            if index == transparent:
-                return True
-    elif img.mode == "RGBA":
-        extrema = img.getextrema()
-        if extrema[3][0] < 255:
-            return True
-    return False
-
 def image_properties(img):
-    global img_mode
     if img:
-        if has_transparency(img):
-            img_mode = "RGBA"
-        else:
-            img_mode = "RGB"
-        return f"Resolution: Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img_mode}"
+        # Chỉ báo thông tin trực tiếp từ ảnh, không kiểm tra alpha
+        return f"Resolution: Width: {img.size[0]}, Height: {img.size[1]}  |  Color Mode: {img.mode}"
 
 def reset():
     global last_file
@@ -184,9 +165,21 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
 
     upsampler = get_upsampler(model_name, denoise_strength)
 
-    # PIL -> cv2 BGRA
+    # PIL -> cv2 (giữ nguyên nếu có alpha; ta sẽ bỏ alpha trước khi lưu JPG)
     cv_img = numpy.array(img)
-    img_bgra = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
+    if cv_img.ndim == 3 and cv_img.shape[2] == 4:
+        # RGBA -> BGRA
+        img_bgra = cv2.cvtColor(cv_img, cv2.COLOR_RGBA2BGRA)
+    elif cv_img.ndim == 3 and cv_img.shape[2] == 3:
+        # RGB -> BGR, rồi thêm alpha giả để pipeline cũ vẫn chạy nếu cần
+        bgr = cv2.cvtColor(cv_img, cv2.COLOR_RGB2BGR)
+        alpha = numpy.full((bgr.shape[0], bgr.shape[1], 1), 255, dtype=bgr.dtype)
+        img_bgra = numpy.concatenate([bgr, alpha], axis=2)
+    else:
+        # 1-channel (L) -> BGR + alpha
+        bgr = cv2.cvtColor(cv_img, cv2.COLOR_GRAY2BGR)
+        alpha = numpy.full((bgr.shape[0], bgr.shape[1], 1), 255, dtype=bgr.dtype)
+        img_bgra = numpy.concatenate([bgr, alpha], axis=2)
 
     try:
         if face_enhance:
@@ -201,9 +194,15 @@ def realesrgan(img, model_name, denoise_strength, face_enhance, outscale):
         print('Error', error)
         return None
     else:
-        extension = 'png' if img_mode == 'RGBA' else 'jpg'
-        out_filename = f"output_{rnd_string(8)}.{extension}"
-        cv2.imwrite(out_filename, output)
+        out_filename = f"output_{rnd_string(8)}.jpg"
+        # Đảm bảo ảnh 3 kênh trước khi lưu JPG
+        if output.ndim == 3 and output.shape[2] == 4:
+            output_to_save = cv2.cvtColor(output, cv2.COLOR_BGRA2BGR)
+        elif output.ndim == 3 and output.shape[2] == 3:
+            output_to_save = output
+        else:
+            output_to_save = cv2.cvtColor(output, cv2.COLOR_GRAY2BGR)
+        cv2.imwrite(out_filename, output_to_save)
         global last_file
         last_file = out_filename
         return out_filename
@@ -237,7 +236,7 @@ def main():
             with gr.Group():
                 input_image = gr.Image(label="Input Image", type="pil", image_mode="RGBA")
                 input_image_properties = gr.Textbox(label="Image Properties", max_lines=1)
-            output_image = gr.Image(label="Output Image", image_mode="RGBA")
+            output_image = gr.Image(label="Output Image", image_mode="RGB")
         with gr.Row():
             reset_btn = gr.Button("Remove images")
             restore_btn = gr.Button("Upscale")