ACAI/model-yolo-person-classify
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

dk - edit

Browse Source
This commit is contained in:
CHO Daekeun
2026-07-02 14:06:02 +09:00
parent 0007eb5709
commit 55fc668058
2 changed files with 139 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main.py
+16 -133
View File
@@ -1,145 +1,28 @@
import os # 시스템 생성 — 이 파일은 커스텀 UI 입력에 따라 자동 생성됩니다. 직접 수정하지 마세요.
import json
import torch
import argparse import argparse
import requests
from io import BytesIO
from PIL import Image
from torchvision import transforms as T
from net import get_model
from clearml import Task # 1. ClearML 임포트 from clearml import Task
num_cls_dict = {'market': 30, 'duke': 23} from start import run
num_ids_dict = {'market': 751, 'duke': 702}
transforms = T.Compose([
T.Resize(size=(288, 144)),
T.ToTensor(),
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
def init_acai_task(output):
task = Task.init(
project_name="Person_Attribute_Recognition",
task_name="model-yolo-person-classify"
)
def upload_acai_task(task: Task, output: dict) -> None:
"""UI output 정의에 맞춰 업로드"""
result_data = {"output": output, "status": "PASS"} result_data = {"output": output, "status": "PASS"}
task.upload_artifact(name="final_result", artifact_object=result_data) task.upload_artifact(name="final_result", artifact_object=result_data)
class PredictDecoder(object):
def __init__(self, dataset):
with open('./doc/label.json', 'r') as f:
self.label_list = json.load(f)[dataset]
with open('./doc/attribute.json', 'r') as f:
self.attribute_dict = json.load(f)[dataset]
self.dataset = dataset
self.num_label = len(self.label_list)
def decode(self, pred):
pred = pred.squeeze(dim=0)
results = []
for idx in range(self.num_label):
name, choice = self.attribute_dict[self.label_list[idx]]
value = choice[pred[idx]]
if value:
results.append((name, value))
return results
def load_network(network, dataset, model_name):
save_path = os.path.join('./checkpoints', dataset, model_name, 'net_last.pth')
network.load_state_dict(torch.load(save_path))
print(f'[+] 모델 로드 완료: {save_path}')
return network
def preprocess_image(image):
src = transforms(image)
return src.unsqueeze(dim=0)
def load_image_from_url_or_path(image_source):
if image_source.startswith(('http://', 'https://')):
response = requests.get(image_source)
response.raise_for_status()
return Image.open(BytesIO(response.content)).convert('RGB')
if not os.path.isfile(image_source):
print(f"Error: Image not found: {image_source}")
exit(1)
return Image.open(image_source).convert('RGB')
def main(image, dataset='market', backbone='resnet50', use_id=False):
assert dataset in ['market', 'duke']
assert backbone in ['resnet50', 'resnet34', 'resnet18', 'densenet121']
model_name = f'{backbone}_nfc_id' if use_id else f'{backbone}_nfc'
num_label = num_cls_dict[dataset]
num_id = num_ids_dict[dataset]
model = get_model(model_name, num_label, use_id=use_id, num_id=num_id)
model = load_network(model, dataset, model_name)
model.eval()
src = preprocess_image(image)
with torch.no_grad():
if not use_id:
out = model.forward(src)
else:
out, _ = model.forward(src)
pred = torch.gt(out, torch.ones_like(out) / 2)
decoder = PredictDecoder(dataset)
results = decoder.decode(pred)
print("\n" + "=" * 50)
print(" Person 상세 속성 분석 결과 ")
print("=" * 50)
for name, value in results:
print(f'{name}: {value}')
print("=" * 50)
init_acai_task(results)
if __name__ == "__main__": if __name__ == "__main__":
# python main.py --image_url "https://acai.ketidev.kr:20443/detect/image/202606/20260619_145116_image.jpg" --xywh "404,290,74,193" task = Task.init(
parser = argparse.ArgumentParser(description="Person Attribute Recognition Agent") project_name="Person_Attribute_Recognition",
parser.add_argument("--image_url", type=str) # --image_url "https://acai.ketidev.kr:20443/detect/image/202606/20260619_145116_image.jpg" task_name="model-yolo-person-classify",
parser.add_argument("--xywh", type=str) # --xywh "x,y,w,h"; (--xywh "404,290,74,193") )
parser = argparse.ArgumentParser()
parser.add_argument("--image_url", type=str, required=True)
parser.add_argument("--xywh", type=str, required=True)
args = parser.parse_args() args = parser.parse_args()
image_url = args.image_url task.connect(parser)
if image_url is None:
print("Error: Image URL or path is required")
exit(1)
oimg = load_image_from_url_or_path(image_url) output = run(args.image_url, args.xywh)
upload_acai_task(task, output)
xywh = args.xywh
if xywh is None:
print("Error: XYWH is required")
exit(1)
xywh = xywh.split(",")
x = int(xywh[0])
y = int(xywh[1])
w = int(xywh[2])
h = int(xywh[3])
cimg = oimg.crop((x, y, x + w, y + h))
cimg.save("cropped_image.jpg")
print(f"[+] 타겟 영역 크롭 완료 -> 'cropped_image.jpg' 저장 완료.")
main(cimg)
start.py
+123
View File
@@ -0,0 +1,123 @@
import os
import json
import torch
import argparse
import requests
from io import BytesIO
from PIL import Image
from torchvision import transforms as T
from net import get_model
num_cls_dict = {'market': 30, 'duke': 23}
num_ids_dict = {'market': 751, 'duke': 702}
transforms = T.Compose([
T.Resize(size=(288, 144)),
T.ToTensor(),
T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
class PredictDecoder(object):
def __init__(self, dataset):
with open('./doc/label.json', 'r') as f:
self.label_list = json.load(f)[dataset]
with open('./doc/attribute.json', 'r') as f:
self.attribute_dict = json.load(f)[dataset]
self.dataset = dataset
self.num_label = len(self.label_list)
def decode(self, pred):
pred = pred.squeeze(dim=0)
results = {}
for idx in range(self.num_label):
name, choice = self.attribute_dict[self.label_list[idx]]
value = choice[pred[idx]]
if value:
results[name] = value
return results
def load_network(network, dataset, model_name):
save_path = os.path.join('./checkpoints', dataset, model_name, 'net_last.pth')
network.load_state_dict(torch.load(save_path))
print(f'[+] 모델 로드 완료: {save_path}')
return network
def preprocess_image(image):
src = transforms(image)
return src.unsqueeze(dim=0)
def load_image_from_url_or_path(image_source):
if image_source.startswith(('http://', 'https://')):
response = requests.get(image_source)
response.raise_for_status()
return Image.open(BytesIO(response.content)).convert('RGB')
if not os.path.isfile(image_source):
raise FileNotFoundError(f"Image not found: {image_source}")
return Image.open(image_source).convert('RGB')
def parse_xywh(xywh):
parts = xywh.split(",")
if len(parts) != 4:
raise ValueError(f"xywh must be 'x,y,w,h', got: {xywh}")
return tuple(int(v) for v in parts)
def predict(image, dataset='market', backbone='resnet50', use_id=False):
assert dataset in ['market', 'duke']
assert backbone in ['resnet50', 'resnet34', 'resnet18', 'densenet121']
model_name = f'{backbone}_nfc_id' if use_id else f'{backbone}_nfc'
num_label = num_cls_dict[dataset]
num_id = num_ids_dict[dataset]
model = get_model(model_name, num_label, use_id=use_id, num_id=num_id)
model = load_network(model, dataset, model_name)
model.eval()
src = preprocess_image(image)
with torch.no_grad():
if not use_id:
out = model.forward(src)
else:
out, _ = model.forward(src)
pred = torch.gt(out, torch.ones_like(out) / 2)
decoder = PredictDecoder(dataset)
return decoder.decode(pred)
def run(image_url, xywh, dataset='market', backbone='resnet50', use_id=False):
oimg = load_image_from_url_or_path(image_url)
x, y, w, h = parse_xywh(xywh)
cimg = oimg.crop((x, y, x + w, y + h))
results = predict(cimg, dataset=dataset, backbone=backbone, use_id=use_id)
print("\n" + "=" * 50)
print(" Person 상세 속성 분석 결과 ")
print("=" * 50)
for name, value in results.items():
print(f'{name}: {value}')
print("=" * 50)
return results
if __name__ == "__main__":
# python start.py --image_url "https://..." --xywh "404,290,74,193"
parser = argparse.ArgumentParser(description="Person Attribute Recognition Agent")
parser.add_argument("--image_url", type=str, required=True)
parser.add_argument("--xywh", type=str, required=True)
args = parser.parse_args()
run(args.image_url, args.xywh)