Update app.py

app.py

@@ -1,168 +1,45 @@
 import gradio as gr
-import aiohttp
-import asyncio
-import time
-from datetime import datetime
-import random
-
-async def make_session_request(url, session_id, request_id):
-    """각 세션별 요청을 처리하는 함수"""
-    async with aiohttp.ClientSession() as session:
-        headers = {
-            'User-Agent': f'Mozilla/5.0 (Windows NT 10.0; Win64; x64) Chrome/{random.randint(70, 100)}.0',
-            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8',
-            'Accept-Language': 'en-US,en;q=0.5',
-            'Connection': 'keep-alive',
-            'DNT': '1',
-            'Upgrade-Insecure-Requests': '1'
-        }
-        start_time = time.time()
-        try:
-            async with session.get(url, headers=headers, ssl=False, timeout=30) as response:
-                await response.read()  # 실제로 콘텐츠를 읽음
-                end_time = time.time()
-                return {
-                    'session_id': session_id,
-                    'request_id': request_id,
-                    'status': response.status,
-                    'time': end_time - start_time,
-                    'timestamp': datetime.now().strftime('%H:%M:%S.%f')[:-3]
-                }
-        except Exception as e:
-            end_time = time.time()
-            return {
-                'session_id': session_id,
-                'request_id': request_id,
-                'status': f"Error: {str(e)}",
-                'time': end_time - start_time,
-                'timestamp': datetime.now().strftime('%H:%M:%S.%f')[:-3]
-            }
-
-async def run_session(url, session_id, num_requests, progress_queue):
-    """각 세션의 요청들을 관리하는 함수"""
-    tasks = []
-    for i in range(num_requests):
-        task = make_session_request(url, session_id, i+1)
-        tasks.append(task)
-        await asyncio.sleep(random.uniform(0.1, 0.5))  # 자연스러운 요청 간격
-    
-    results = await asyncio.gather(*tasks)
-    await progress_queue.put((session_id, len(results)))
-    return results
-
-async def load_test(url, num_sessions, requests_per_session, progress=gr.Progress()):
-    """메인 테스트 함수"""
-    if not url.startswith(('http://', 'https://')):
-        url = 'http://' + url
-
-    start_time = time.time()
-    progress_queue = asyncio.Queue()
-    total_requests = num_sessions * requests_per_session
-    completed_requests = 0
-
-    # 프로그레스 모니터링을 위한 비동기 함수
-    async def monitor_progress():
-        nonlocal completed_requests
-        while completed_requests < total_requests:
-            session_id, count = await progress_queue.get()
-            completed_requests += count
-            progress(completed_requests / total_requests, 
-                    desc=f"진행 중... ({completed_requests}/{total_requests} requests)")
-
-    # 모든 세션 시작
-    tasks = []
-    for session_id in range(num_sessions):
-        task = run_session(url, session_id + 1, requests_per_session, progress_queue)
-        tasks.append(task)
-    
-    # 프로그레스 모니터링 태스크 시작
-    monitor_task = asyncio.create_task(monitor_progress())
-    
-    # 모든 세션의 결과 수집
-    all_results = []
-    for session_results in await asyncio.gather(*tasks):
-        all_results.extend(session_results)
-
-    await monitor_task
-    
-    end_time = time.time()
-    total_time = end_time - start_time
-
-    # 결과 분석
-    success_count = sum(1 for r in all_results if isinstance(r['status'], int) and 200 <= r['status'] < 300)
-    failed_count = len(all_results) - success_count
-    response_times = [r['time'] for r in all_results]
-
-    result_text = f"""테스트 결과:
-======================
-총 세션 수: {num_sessions}
-세션당 요청 수: {requests_per_session}
-총 요청 수: {total_requests}
-성공: {success_count}
-실패: {failed_count}
-총 소요 시간: {total_time:.2f}초
-초당 요청 수: {total_requests/total_time:.1f}
-
-응답 시간 통계:
-- 평균: {sum(response_times)/len(response_times):.3f}초
-- 최소: {min(response_times):.3f}초
-- 최대: {max(response_times):.3f}초
-
-상세 요청 로그:
-======================
+import utils
+
+def choose_encode(inp_im, inp_mark, cho):
+    if cho == "stegan":
+        out_im, out_msg = utils.encode(inp_im, inp_mark)
+        return out_im, out_msg
+    if cho == "pnginfo":
+        out_im, out_msg = utils.png_encode(inp_im, inp_mark)
+        return out_im, out_msg
+
+css = """
+footer {
+    visibility: hidden;
+}
 """
-    
-    all_results.sort(key=lambda x: x['timestamp'])
-    for r in all_results[:50]:  # 처음 50개 요청만 표시
-        result_text += f"[{r['timestamp']}] Session {r['session_id']:2d} - Request {r['request_id']:3d} - Status: {r['status']} - Time: {r['time']:.3f}s\n"
-    
-    if len(all_results) > 50:
-        result_text += f"\n... 및 {len(all_results) - 50}개의 추가 요청"
-
-    return result_text
-
-def sync_load_test(url, num_sessions, requests_per_session, progress=gr.Progress()):
-    return asyncio.run(load_test(url, int(num_sessions), int(requests_per_session), progress))
-
-with gr.Blocks() as demo:
-    gr.Markdown("# 동시 접속 부하 테스트")
-    
-    with gr.Row():
-        url_input = gr.Textbox(
-            label="대상 URL",
-            placeholder="example.com or http://example.com"
-        )
-    
-    with gr.Row():
-        session_count = gr.Number(
-            label="동시 세션 수",
-            value=10,
-            minimum=1,
-            maximum=100,
-            step=1
-        )
-        requests_per_session = gr.Number(
-            label="세션당 요청 수",
-            value=10,
-            minimum=1,
-            maximum=100,
-            step=1
-        )
-    
-    with gr.Row():
-        test_button = gr.Button("테스트 시작", variant="primary")
-    
-    with gr.Row():
-        output = gr.Textbox(
-            label="테스트 결과",
-            lines=30
-        )
-    
-    test_button.click(
-        fn=sync_load_test,
-        inputs=[url_input, session_count, requests_per_session],
-        outputs=output
-    )
 
-demo.queue()
-demo.launch()
+with gr.Blocks(css=css) as app:
+    gr.HTML("""<a href="https://visitorbadge.io/status?path=https%3A%2F%2Ffantos-watermark.hf.space">
+               <img src="https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Ffantos-watermark.hf.space&countColor=%23263759" />
+               </a>""")
+    
+    with gr.Tab("Add Watermark"):
+        cho = gr.Radio(choices=["stegan", "pnginfo"], value="stegan")
+        with gr.Row():
+            with gr.Column():
+                inp_im = gr.Image(label="Input Image", type="filepath")
+                inp_mark = gr.Textbox(label="Watermark")
+                mark_btn = gr.Button("Add Watermark")
+                msg_box = gr.Textbox(label="System Message")
+            with gr.Column():
+                out_im = gr.Image(label="Watermarked Image")
+    
+    with gr.Tab("Detect Watermark"):
+        with gr.Row():
+            with gr.Column():
+                det_im = gr.Image(label="Watermarked Image", type="filepath")
+                det_btn = gr.Button("Detect")
+            with gr.Column():
+                det_msg = gr.Textbox(label="Detected Watermark", lines=6, max_lines=50)
+    
+    mark_btn.click(choose_encode, [inp_im, inp_mark, cho], [out_im, msg_box])
+    det_btn.click(utils.decode, [det_im], det_msg)
+
+app.launch()