Spaces:
Sleeping
Sleeping
File size: 3,031 Bytes
e3421f8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
# utils.py
import numpy as np
from pathlib import Path
from typing import List, Tuple, Optional
import torch
import config
def get_available_models() -> List[str]:
"""Lấy danh sách tất cả model trong folder models"""
models_dir = Path(config.CONFIG['models_dir'])
model_files = list(models_dir.glob("*.pth"))
return sorted([f.name for f in model_files])
def validate_model_file(model_path: str) -> bool:
"""Kiểm tra xem file model có hợp lệ không"""
path = Path(model_path)
if not path.exists():
return False
if path.suffix != '.pth':
return False
if path.stat().st_size < 1000: # File quá nhỏ
return False
return True
def format_time(seconds: float) -> str:
"""Format thời gian từ giây sang HH:MM:SS"""
hours = int(seconds // 3600)
minutes = int((seconds % 3600) // 60)
secs = int(seconds % 60)
if hours > 0:
return f"{hours}h {minutes}m {secs}s"
elif minutes > 0:
return f"{minutes}m {secs}s"
else:
return f"{secs}s"
def get_color_for_value(color_val: int) -> str:
"""Lấy màu hex từ giá trị color"""
return config.COLORS_MAP.get(color_val, '#ffffff')
def create_bottle_html(bottle_state: np.ndarray, bottle_idx: int,
selected: bool = False) -> str:
"""Tạo HTML cho một chai"""
border_color = '#FFD700' if selected else '#333'
border_width = '3' if selected else '2'
html = f'<div style="text-align: center; margin: 10px;">'
html += f'<div style="width: 60px; height: 150px; border: {border_width}px solid {border_color}; '
html += f'background: #f0f0f0; margin-bottom: 10px; display: flex; flex-direction: column-reverse; overflow: hidden;">'
for height_idx in range(len(bottle_state)):
color_val = int(bottle_state[height_idx])
color = get_color_for_value(color_val)
html += f'<div style="width: 100%; height: 30px; background: {color}; border-bottom: 1px solid #999;"></div>'
html += '</div>'
html += f'<p style="margin: 5px 0; font-weight: bold;">Chai {bottle_idx}</p>'
html += '</div>'
return html
def get_device() -> torch.device:
"""Lấy device (GPU hoặc CPU)"""
device_config = config.MODEL_CONFIG['device']
if device_config == 'auto':
return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
elif device_config == 'cuda':
if torch.cuda.is_available():
return torch.device('cuda')
else:
print("Warning: CUDA not available, falling back to CPU")
return torch.device('cpu')
else:
return torch.device('cpu')
def print_device_info():
"""In thông tin về device"""
device = get_device()
print(f"Device: {device}")
if device.type == 'cuda':
print(f"GPU Name: {torch.cuda.get_device_name(0)}")
print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB") |